﻿using ESPCS;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using ESPCS;

using Microsoft.WindowsAPICodePack.Shell;
using Microsoft.WindowsAPICodePack.Shell.PropertySystem;

namespace AZPEP_SysRi_Tool
{
    public class DiskPart
    {
        /// <summary>
        /// DiskPart文件路径
        /// </summary>
        private static readonly string DiskPartPath = Environment.SystemDirectory + @"\diskpart.exe";

        /// <summary>
        /// DiskPart结果输出目录（有可能不需要）
        /// </summary>
        private static readonly string DiskPartResultsPath = Path.GetTempPath() + @"diskpart.tmp";

        /// <summary>
        /// 自动化创建分区
        /// </summary>
        /// <param name="SystemPartitionLetter">系统分区盘符</param>
        /// <param name="SystemSizeMB">系统大小（MB）</param>
        /// 
        /// <param name="SizeMB">分区大小（MB）</param>
        /// <param name="Letter">盘符</param>
        /// <param name="ShrinkPartitionLetter">压缩的卷，大概率不需要</param>
        /// <returns></returns>
        /// 
        public static bool CreatePEDisk(
            char SystemPartitionLetter,
            long SystemSizeMB,

            long SizeMB,

            out char? Letter,
            out char? ShrinkPartitionLetter,
            out ErrCode.ErrorCode Error
            )
        {
            /* 获取多余盘符 */
            List<char> TotalLetter = new List<char> { 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z' };
            List<char> AllLetter = new List<char>();

            foreach (Disk.PartitionInfo PI in Disk.GetAllPartition())
            {
                AllLetter.Add(PI.partitionLetter[0]);
            }

            List<char> FreeLetter = new List<char>();
            foreach (char L in TotalLetter)
            {
                if (AllLetter.IndexOf(L) == -1)
                {
                    FreeLetter.Add(L);
                }
            }

            //如果没有获取到多余的盘符，报错NO_EXTRA_DRIVE_LETTERS并返回
            if (FreeLetter.Count == 0)
            {
                MyError.ShowError(ErrCode.ErrorCode.__MAINPROG__CREATEPE_NO_MORE_DRIVE_LETTERS, "");
                Letter = null;
                ShrinkPartitionLetter = null;
                Error = ErrCode.ErrorCode.__MAINPROG__CREATEPE_NO_MORE_DRIVE_LETTERS;
                return false;
            }

            // 设置PE释放盘符
            char FinallyLetter = FreeLetter[0];


            /* 计算大小 */
            //decimal SizeMB = SizeByte;      // Byte换算为MB(1 MB = 1000 Bytes)
            decimal SizeMB_up = Math.Ceiling((decimal)SizeMB);   //向上取整

            //对齐到128的整数倍
            decimal FinallySizeMB = SizeMB_up + (SizeMB_up % 128);


            LogProtocol.WriteLog.Debug($"最终创建的分区大小：{FinallySizeMB} MB");

            //获取硬盘数
            int HarddiskNumber = ESPCS.ESP.GetHarddiskNum();

            for (int i = 0; i < HarddiskNumber; i++)
            {
                //尝试创建分区
                if (
                CreatePartition(
                    0,
                    FinallyLetter,
                    FinallySizeMB
                    )
                )
                {
                    Letter = FinallyLetter;
                    Error = ErrCode.ErrorCode.SUCCESS;
                    ShrinkPartitionLetter = null;
                    return true;
                }
            }

            /* 创建不成功，尝试压缩分区后再创建分区 */

            foreach (char L in AllLetter)
            {
                //是否开启Bitlocker，开启就跳过
                if (Bitlocker.CheckVolumeBitlocker(L)) continue;

                //获取分区信息
                DriveInfo DI = DiskInfo.GetVolumeInfo(L);

                //分区未准备好或者分区不是固定磁盘，跳过
                if (DI.IsReady == false || DI.DriveType != DriveType.Fixed) continue;

                //如果分区为系统盘，且压缩后的分区大小不足释放后映像大小的1.5倍，跳过
                if (L.ToString().ToLower() == SystemPartitionLetter.ToString().ToLower())
                {
                    if (DI.TotalSize - SizeMB * 1024 < SystemSizeMB * 1024 * 1.5) continue;
                }

                //压缩分区
                if (ShrinkPartition(L, FinallySizeMB + 24))
                {
                    for (int i = 0; i < HarddiskNumber; i++)
                    {
                        if (CreatePartition(
                            0,
                            FinallyLetter,
                            FinallySizeMB
                            )
                        )
                        {
                            Letter = FinallyLetter;
                            Error = Error = ErrCode.ErrorCode.SUCCESS;
                            ShrinkPartitionLetter = L;
                            return true;
                        }
                    }
                }
            }

            /* 再不成功，报错返回 */
            MyError.ShowError(ErrCode.ErrorCode.__MAINPROG__CREATEPE_NO_MORE_DRIVE_LETTERS, "");
            Letter = null;
            Error = ErrCode.ErrorCode.__MAINPROG__CREATEPE__DISKPART__CREATE_PARTITION_FAIL;
            ShrinkPartitionLetter = null;
            return false;


        }

        /// <summary>
        /// 删除卷
        /// </summary>
        /// <param name="VolumeLetter">卷盘符</param>
        /// <returns></returns>
        public static bool DeleteVolume(
            char VolumeLetter
            )
        {
            string Script =
                $"SEL VOLUME={VolumeLetter}\n" +
                $"DELETE VOLUME\n";

            using (StreamWriter sw = new StreamWriter($"{System.Environment.GetEnvironmentVariable("TEMP")}\\DiskScript.script", false))
            {
                sw.Write(Script);
            }

            int ExitCode = 0;
            string Result;

            RunDiskPart(
                $"{System.Environment.GetEnvironmentVariable("TEMP")}\\DiskScript.script",
                out ExitCode,
                out Result
                );

            if (ExitCode != 0) { LogProtocol.WriteLog.Debug($"\n{Result}"); }

            if (ExitCode == 0) { return true; }
            else { return false; }
        }
        /// <summary>
        /// 扩容卷
        /// </summary>
        /// <param name="VolumeLetter">卷盘符</param>
        /// <returns></returns>
        public static bool ExtendVolume
            (
            char VolumeLetter
            )
        {
            string Script =
                $"SEL VOLUME={VolumeLetter}\n" +
                $"EXTEND\n";

            using (StreamWriter sw = new StreamWriter($"{System.Environment.GetEnvironmentVariable("TEMP")}\\DiskScript.script", false))
            {
                sw.Write(Script);
            }

            int ExitCode = 0;
            string Result;

            RunDiskPart(
                $"{System.Environment.GetEnvironmentVariable("TEMP")}\\DiskScript.script",
                out ExitCode,
                out Result
                );

            if (ExitCode != 0) { LogProtocol.WriteLog.Debug($"\n{Result}"); }

            if (ExitCode == 0) { return true; }
            else { return false; }
        }
        /// <summary>
        /// 压缩分区（压缩卷）
        /// </summary>
        /// <param name="VolumeLetter">卷盘符</param>
        /// <param name="ShrinkSizeMB">压缩大小（MB）</param>
        /// <returns></returns>
        public static bool ShrinkPartition(
            in char VolumeLetter,
            in decimal ShrinkSizeMB
            )
        {
            string Script =
                $"SEL VOLUME={VolumeLetter}\n" +
                $"SHRINK DESIRED = {ShrinkSizeMB}\n";

            using (StreamWriter sw = new StreamWriter($"{System.Environment.GetEnvironmentVariable("TEMP")}\\DiskScript.script", false))
            {
                sw.Write(Script);
            }

            int ExitCode = 0;
            string Result;

            RunDiskPart(
                $"{System.Environment.GetEnvironmentVariable("TEMP")}\\DiskScript.script",
                out ExitCode,
                out Result
                );

            if (ExitCode != 0) { LogProtocol.WriteLog.Debug($"\n{Result}"); }
            if (ExitCode == 0) { return true; }
            else { return false; }
        }
        /// <summary>
        /// 创建分区
        /// </summary>
        /// <param name="DiskNumber">硬盘号</param>
        /// <param name="Letter">盘符，应为"C"而不是"C:\"</param>
        /// <param name="Size">大小(MB)</param>
        /// <returns></returns>
        public static bool CreatePartition(
            in int DiskNumber,
            in char Letter,
            in decimal SizeMB
            )
        {
            string Script =
                $"SEL DISK={DiskNumber}\n" +
                $"CREATE PARTITION PRIMARY SIZE={SizeMB}\n" +
                $"FORMAT QUICK\n" +
                $"ASSIGN LETTER={Letter}\n";

            using (StreamWriter sw = new StreamWriter($"{System.Environment.GetEnvironmentVariable("TEMP")}\\DiskScript.script", false))
            {
                sw.Write(Script);
            }

            int ExitCode = 0;
            string Result;

            RunDiskPart(
                $"{System.Environment.GetEnvironmentVariable("TEMP")}\\DiskScript.script",
                out ExitCode,
                out Result
                );

            if (ExitCode != 0) { LogProtocol.WriteLog.Debug($"\n{Result}"); }

            if (ExitCode == 0) { return true; }
            else { return false; }


        }


        /// <summary>
        /// 运行DiskPart，执行脚本
        /// </summary>
        /// <param name="ScriptFilePath">脚本路径</param>
        /// <param name="ExitCode">退出代码</param>
        /// <param name="Result">返回结果</param>
        private static bool RunDiskPart(
            in string ScriptFilePath,

            out int ExitCode,
            out string Result
            )
        {

            Process p = new Process();
            //设置要启动的应用程序
            p.StartInfo.FileName = DiskPartPath;
            //是否使用操作系统shell启动
            p.StartInfo.UseShellExecute = false;
            // 接受来自调用程序的输入信息
            p.StartInfo.RedirectStandardInput = true;
            //输出信息
            p.StartInfo.RedirectStandardOutput = true;
            // 输出错误
            p.StartInfo.RedirectStandardError = true;
            //不显示程序窗口
            p.StartInfo.CreateNoWindow = true;

            //p.StartInfo.Arguments = $"/s \"{ScriptFilePath}\" > \"{DiskPartResultsPath}\"";
            p.StartInfo.Arguments = $"/s \"{ScriptFilePath}\"";
            p.Start();

            string Text;
            Text = p.StandardOutput.ReadToEnd();

            p.WaitForExit();


            ExitCode = p.ExitCode;
            Result = Text;

            LogProtocol.WriteLog.Debug($"DiskPart执行完毕，退出代码{ExitCode}。\n返回内容：\n{Result}");

            if (ExitCode == 0) { return true; }
            else { return false; }

        }
    }

    public class Disk
    {
        /// <summary>
        /// 分区信息结构体
        /// </summary>
        public struct PartitionInfo
        {
            /// <summary>
            /// 盘符，如"C:\"
            /// </summary>
            public string partitionLetter;

            /// <summary>
            /// 驱动器类型
            /// </summary>
            public System.IO.DriveType partitionType;

            /// <summary>
            /// 驱动器名称
            /// </summary>
            public string partitionName;

            /// <summary>
            /// 驱动器总大小
            /// </summary>
            public long partitionTotalSpace;

            /// <summary>
            /// 驱动器可用空间大小
            /// </summary>
            public long partitionFreeSpace;
        }

        /// <summary>
        /// 获取计算机上所有已挂载的盘符
        /// </summary>
        /// <returns></returns>
        public static List<PartitionInfo> GetAllPartition()
        {
            //获取所有磁盘基本信息
            DriveInfo[] DIs = System.IO.DriveInfo.GetDrives();

            //遍历数组
            List<PartitionInfo> PIL = new List<PartitionInfo>();

            //设置信息
            foreach (System.IO.DriveInfo DI in DIs)
            {
                PartitionInfo info = new PartitionInfo();
                info.partitionLetter = DI.Name;
                info.partitionType = DI.DriveType;
                info.partitionName = DI.VolumeLabel;
                info.partitionTotalSpace = DI.TotalSize;
                info.partitionFreeSpace = DI.TotalFreeSpace;

                PIL.Add(info);

            }

            return PIL;
        }

        /// <summary>
        /// 获取单个分区的信息
        /// </summary>
        /// <param name="letter">盘符，如"C:\"</param>
        /// <returns>成功返回PartitionInfo，失败返回null</returns>
        public static PartitionInfo GetPartitionInfo(string letter)
        {
            List<PartitionInfo> PIs = GetAllPartition();

            foreach (PartitionInfo PI in PIs)
            {
                if (PI.partitionLetter == letter) return PI;
            }

            throw new Exception("分区未找到！");
        }

    }

    public class Bitlocker
    {
        /// <summary>
        /// 检查卷的Bitlocker是否开启
        /// </summary>
        /// <param name="Letter">盘符</param>
        /// <returns>开启返回true，未开启返回false</returns>
        public static bool CheckVolumeBitlocker(char Letter)
        {
            IShellProperty prop = ShellObject.FromParsingName($"{Letter}:").Properties.GetProperty("System.Volume.BitLockerProtection");
            int? bitLockerProtectionStatus = (prop as ShellProperty<int?>).Value;

            if (bitLockerProtectionStatus.HasValue && (bitLockerProtectionStatus == 1 || bitLockerProtectionStatus == 3 || bitLockerProtectionStatus == 5)) return true;
            else return false;
        }
    }

    public class DiskInfo
    {
        public static DriveInfo GetVolumeInfo(char Letter)
        {
            DriveInfo[] allDirves = DriveInfo.GetDrives();

            foreach (DriveInfo DI in allDirves)
            {
                if (DI.Name.ToLower() == Letter.ToString().ToLower()) return DI;
            }

            MyError.ShowError(ErrCode.ErrorCode.WHAT_THE_FXXK, "驱动器“{0}”之前存在，现在不存在？搞什么鬼？");
            MyError.EmergencyExitEx(1);

            return null;
        }
    }
}
